DEBUG = False
import random
import pprint
try:
    import cPickle as pickle
except ImportError:
    import pickle

try:
    import numpy
    numpy.seterr(all='warn', over='raise')
except ImportError:
    pass
#     
# try:
#     import psyco
# except ImportError:
#     print 'no psyco'
#     pass

class MarkovChain(object):
    def __init__(self, states, start_probabilities, state_transitions):
        observations_prob = {}
        for state in states:
            observations_prob[state] = {state:1}
            
        self.hmm = MyHMM(states, states, start_probabilities, state_transitions, observations_prob)

    def train(self,state_sequences, normalize = False):
        self.hmm.train(state_sequences, [], normalize)
        observations_prob = {}
        for state in self.hmm.states:
            observations_prob[state] = {state:1}
        self.hmm.observations_prob=observations_prob
        self.hmm.observations = self.hmm.states
    
    def probability(self, sequence):
        return self.hmm.forward(sequence)
    

            
class MyHMM(object):
    def __init__(self, states, observations, start_probabilities, state_transitions, observations_prob):
        """coef is used to allow for unseen data. Use the probability of a word that has one or two occurences"""
        self.states = states
        self.observations = observations
        self.start_probabilities = start_probabilities
        self.state_transitions = state_transitions
        self.observations_prob = observations_prob
        self.len_states = len(self.states)
        self.len_observations = len(self.observations)
        self._optimized=False
        self._numpy = True
    
    @property
    def usesNumpy(self):
        return self._numpy and self._optimized
    
    @property
    def coef_trans(self):
        return self._coef #1e-150
    
    @property
    def coef_obs(self):
        """docstring for coef_obs"""
        #dividing by 100.0 makes the unit tests pass
        #the concept is that we are less likely not to have found an observation than a state transition.
        return self._coef/100.0 #1e-180
    
    def _calc_coef(self):
        """Calculate the smoothing coef, which is the smallest probability divided by 100.0"""
        min_prob = 1.0
        for state in self.state_transitions.values():
            if min_prob>min(state.values()):
                min_prob=min(state.values())
        for state in self.observations_prob.values():
            if min_prob>min(state.values()):
                min_prob=min(state.values())
        self._coef = min_prob/100.0
        # print 'self._coef',self._coef
    
    @classmethod
    def read_model(cls, datafile):
        "read a model from a file"
        global DEBUG
        
        try:
            if DEBUG: print "opening pickle file", datafile
            file =  open(datafile, 'rb')
            if DEBUG: print "unpickling..."
            model = pickle.load(file)
        except:
            return None
        if DEBUG: print "done."
        return model
    
    @classmethod        
    def write_model(cls, model,datafile):
        file =  open(datafile, 'wb')
        model = pickle.dump(model,file,pickle.HIGHEST_PROTOCOL)
    
    def __str__(self):
        if self._optimized: 
            prefix='Optimized,'
            if self._numpy: prefix+=' numpy enabled '
            else: prefix+=' non-numpy '
        else: prefix = 'Non-optimized '
        return prefix+'MyHMM with %d states and %d observations' %(self.len_states, len(self.observations))
    
    def train(self, state_sequences, state_observations, normalize = False): 
        self.train_states(state_sequences, normalize)
        self.train_observations(state_observations, normalize)
    
    def update_self(self):
        self.len_states = len(self.states)
        self.len_observations = len(self.observations)
        self._calc_coef()
    
    def train_states(self, state_sequences, normalize):
        state_count = 0      
        start_count = 0
        for sequence,prob in state_sequences:
            prev = None
            for part in sequence:
                if not part in self.states: self.states.append(part)
                if prev:
                    self.state_transitions[prev] = self.state_transitions.get(prev, {})
                    self.state_transitions[prev][part] = self.state_transitions[prev].get(part, 0) + prob  
                    # state_count += prob
                    if not normalize:
                        assert(self.state_transitions[prev][part]<=1)              
                else:
                    self.start_probabilities[part] = self.start_probabilities.get(part, 0) + prob
                    start_count += prob
                    if not normalize:
                        assert(self.start_probabilities[part]<=1)
                prev = part
                
        if normalize:
            for key in self.state_transitions:
                rowcount = 0
                for key2 in self.state_transitions[key]:
                    rowcount +=self.state_transitions[key][key2]
                for key2 in self.state_transitions[key]:                
                    self.state_transitions[key][key2] = float(self.state_transitions[key][key2])/rowcount
                    assert(self.state_transitions[key][key2]<=1) 
            for key in self.start_probabilities:
                self.start_probabilities[key]= float(self.start_probabilities[key])/start_count
                assert(self.start_probabilities[key]<=1) 
        self.update_self()
    
    def train_observations(self, state_observations, normalize, addStates = False, addProb=False):
        # obs_count = 0
        for sequence,prob in state_observations:
            for (observation, state) in sequence:
                if not state in self.states:
                    if not addStates: print state.encode('utf-8'), 'NOT FOUND IN STATES!!!!!!!!'
                    else: self.states.append(state)
                if not observation in self.observations: self.observations.append(observation)
                self.observations_prob[state] = self.observations_prob.get(state, {})
                if addProb:
                    self.observations_prob[state][observation] = self.observations_prob[state].get(observation,0) + prob
                else:
                    self.observations_prob[state][observation] = prob
                # obs_count += prob
                if not normalize:
                        assert(self.observations_prob[state][observation]<=1)
        if normalize:
                for key in self.observations_prob:
                    rowcount = 0.0
                    for key2 in self.observations_prob[key]:
                        rowcount+= self.observations_prob[key][key2]
                    for key2 in self.observations_prob[key]:
                        self.observations_prob[key][key2] = float(self.observations_prob[key][key2])/rowcount
                        assert(self.observations_prob[key][key2]<=1) 
        self.update_self()    
    
    def _optimize_numpy(self):
        """docstring for _optimize_numpy"""
        # print 'optimize numpy'
        
        print 'coef for opt_numpy', self._coef
        _opt_numpy_state_trans = numpy.zeros((self.len_states, self.len_states),dtype='float')
        _opt_numpy_start_prob = numpy.zeros((self.len_states,), dtype = 'float')
        _opt_numpy_obs_prob = numpy.zeros((self.len_states, self.len_observations), dtype = 'float')
        # scale = 1e10
        for i, state_i in enumerate(self.states):
            try:
                a = self.start_probabilities[state_i]
            except KeyError:
                a = self.coef_trans
            _opt_numpy_start_prob[i]=a
            for j, state_j in enumerate(self.states):
                try:
                    b = self.state_transitions[state_i][state_j]
                except KeyError:
                    b = self.coef_trans
                _opt_numpy_state_trans[i,j] = b
            for j, obs_j in enumerate(self.observations):
                try:
                    c = self.observations_prob[state_i][obs_j]
                except KeyError:
                    c = self.coef_obs
                _opt_numpy_obs_prob[i,j]= c
        
        self._opt_numpy_obs_prob = _opt_numpy_obs_prob
        self._opt_numpy_state_trans = _opt_numpy_state_trans
        self._opt_numpy_start_prob = _opt_numpy_start_prob
    
    def _optimize(self):
        if self._optimized: 
            return
        
        self._calc_coef()
        _opt_start_prob = []
        _opt_obs_prob = []
        _opt_state_trans = []
        
        for i,state_i in enumerate(self.states):
            try:
                a = self.start_probabilities[state_i]
            except KeyError:
                a = self.coef_trans
            _opt_start_prob.append(a)
            _opt_state_trans.append([])
            _opt_obs_prob.append([])
            for state_j in self.states:
                try:
                    b = self.state_transitions[state_i][state_j]
                except KeyError:
                    b = self.coef_trans
                _opt_state_trans[i].append(b)
            for obs_j in self.observations:
                try:
                    c = self.observations_prob[state_i][obs_j]
                except KeyError:
                    c = self.coef_obs
                _opt_obs_prob[i].append(c)

        if self._numpy:
            try:
                import numpy
                self._optimize_numpy()
            except ImportError:
                print 'Using slow mode. Please visit http://numpy.scipy.org and download NumPy.'
        self._opt_obs_prob = _opt_obs_prob
        self._opt_state_trans = _opt_state_trans
        self._opt_start_prob = _opt_start_prob
        self._optimized = True
    
    def viterbi(self,sequence):
        if self._optimized and self._numpy:
            # print 'numpy vit'
            return self._opt_viterbi_numpy(sequence)
        elif self._optimized and not self._numpy:
            # print 'opti vit'
            return self._opt_viterbi(sequence)
        else:
            # print 'stupid vit'
            return self._viterbi(sequence)
    
    def _opt_viterbi_numpy2(self, sequence):
        """docstring for _opt_viterbi_numpy2"""
        
        print '_opt_viterbi_numpy2'
        start_probabilities = self._opt_numpy_start_prob
        observations_prob = self._opt_numpy_obs_prob
        state_transitions = self._opt_numpy_state_trans
        len_seq = len(sequence)
        
        # d = numpy.zeros([len_seq,len(self.states)], dtype='double', order='C')
        # f = numpy.zeros([len_seq,len(self.states)], dtype='int', order='C')
        
        obs_indices = self.indices_from_symbols(sequence)
        
        #construct d,f: d holds values, f holds indices
        d = []
        f = []
        for i in range(0, len(sequence)):
            d.append([])
            f.append([])
        
        index_0 = obs_indices[0]
        for state_i in xrange(len(self.states)):
            c = observations_prob[state_i][index_0]
            d[0].append(start_probabilities[state_i]*c)
        
        if len(sequence)>1:
            for t in xrange(1, len(sequence)): # for each t > 1
                state_i = None
                sequence_t = sequence[t]
                index_t = obs_indices[t]
                for state_i in xrange(len(self.states)):
                    dtemp = []
                    ftemp = []
                    e = observations_prob[state_i][index_t]
                    for state_j in xrange(len(self.states)):
                        daji = d[t-1][state_j] * state_transitions[state_j][state_i]       
                        
                        dtemp.append(daji * e)
                        ftemp.append(daji)
                    # print 'd',
                    # pprint.pprint(dtemp)
                    # print 'f',
                    # pprint.pprint(ftemp)
                    f[t].append(ftemp.index(max(ftemp)))
                    d[t].append(max(dtemp))
        else:
            t=0
            
        # print 'final f:'    
        # pprint.pprint(f)
        # print 'final d:'
        # pprint.pprint(d)
        # print d
        state_seq = []
        it = d[t].index(max(d[t]))
        state_seq.append(self.states[it])
        for t in reversed(range(0, len(sequence)-1)):
            it = f[t+1][it]
            state_seq.append(self.states[it])

        state_seq.reverse()
        return state_seq
    
    def indices_from_symbols(self, sequence):
        """docstring for indices_from_symbols"""
        return [self.observations.index(t) for t in sequence]
    
    def _opt_viterbi_numpy(self, sequence):
        #assume that the instance variables:
        #start_probabilities,observations_prob,state_transitions
        #are now of type numpy.array
        #and we are working only with indexes
        # also assume that there are no blank spots - this can be done by running regular viterbi once
        # print "_opt_viterbi_numpy"
        
        start_probabilities = self._opt_numpy_start_prob
        observations_prob = self._opt_numpy_obs_prob
        state_transitions = self._opt_numpy_state_trans
        len_seq = len(sequence)
        #construct d,f: d holds values, f holds indices
        d = numpy.zeros([len_seq,len(self.states)], dtype='double', order='C')
        f = numpy.zeros([len_seq,len(self.states)], dtype='int', order='C')
        
        obs_indices = self.indices_from_symbols(sequence)

        assert state_transitions.shape == (len(self.states),len(self.states)) 
        assert observations_prob.shape == (len(self.states),len(self.observations))
        assert start_probabilities.shape == (len(self.states),)
        #initialize the start probabilities
        # index_0 = obs_indices[0]
        #these are matrix multiplications
        d[0,:] = start_probabilities * observations_prob[:, obs_indices[0]]
        for t in xrange(1, len_seq): # for each t > 1
            #matrix multiplications
            #the previous states
            daji = d[t-1,:] * state_transitions.T
            #array.max(axis=0) returns the max of each of the columns
            d[t,:] = daji.max(1) * observations_prob[:, obs_indices[t]]
            f[t,:] = daji.argmax(1)
        
        it = numpy.argmax(d[len_seq-1,:])
        state_seq = [self.states[it]]
        for t in reversed(xrange(0, len_seq-1)):
            it = f[t+1,it]
            state_seq.append(self.states[it])

        state_seq.reverse()
        return state_seq
        
    def _opt_viterbi(self, sequence):
        #WARNING DO NOT EDIT THIS METHOD EDIT _opt_viterbi_numpy INSTEAD.
        #THEN COPY PASTE OVER THIS AND CHANGE array[x,y] to array[x][y]!
        #assume that the instance variables:
        #start_probabilities,observations_prob,state_transitions
        #are now of type numpy.array
        #and we are working only with indexes
        # also assume that there are no blank spots - this can be done by running regular viterbi once
        # print "opt_viterbi"
        
        start_probabilities = self._opt_start_prob
        observations_prob = self._opt_obs_prob
        state_transitions = self._opt_state_trans
        
        #construct d,f: d holds values, f holds indices
        d = []
        f = []
        for i in range(0, len(sequence)):
            d.append([])
            f.append([])
        
        obs_indices = []
        for t in sequence:
            try:
                obs_indices.append(self.observations.index(t))
            except IndexError:
                obs_indices.append(None)
        # print obs_indices
        if obs_indices[0] is not None:
            index_0 = obs_indices[0]
            for state_i in xrange(len(self.states)):
                c = observations_prob[state_i][index_0]
                d[0].append(start_probabilities[state_i]*c)
        else:
            d[0] = [coef*start_probabilities[i] for i in xrange(len(self.states))]
        
        if len(sequence)>1:
            for t in xrange(1, len(sequence)): # for each t > 1
                state_i = None
                sequence_t = sequence[t]
                index_t = obs_indices[t]
                for state_i in xrange(len(self.states)):
                    dtemp = []
                    ftemp = []
                    if index_t is not None:
                        e = observations_prob[state_i][index_t]
                        for state_j in xrange(len(self.states)):
                            b = state_transitions[state_j][state_i]                        
                            daji = d[t-1][state_j] * b
                            
                            dtemp.append(daji * e)
                            ftemp.append(daji)
                    else:
                        for state_j in xrange(len(self.states)):
                            b = state_transitions[state_j][state_i]                        
                            daji = d[t-1][state_j] * b                         
                            dtemp.append(daji * coef)
                            ftemp.append(daji)
                        
                    f[t].append(ftemp.index(max(ftemp)))
                    d[t].append(max(dtemp))
        else:
            t=0
        
        # print d
        state_seq = []
        it = d[t].index(max(d[t]))
        state_seq.append(self.states[it])
        for t in reversed(range(0, len(sequence)-1)):
            it = f[t+1][it]
            state_seq.append(self.states[it])

        state_seq.reverse()
        return state_seq         
                   
    def _viterbi(self, sequence):
        global DEBUG
        # print '_viterbi'
        d = []
        f = []
        for i in range(0, len(sequence)):
            d.append([])
            f.append([])
        sequence_0=sequence[0]
        for state_i in self.states:
#            state_i=self.states[i]
            try:
                a = self.start_probabilities[state_i]
            except KeyError:
                a = self.coef_trans
                self.start_probabilities[state_i] = self.coef_trans
                if DEBUG: print 'Start probability not found:',state_i.encode('utf-8')
            try:
                b = self.observations_prob[state_i]
            except KeyError:
                b= {}
                self.observations_prob[state_i] = b
                if DEBUG: 'Observation for state not found:', state_i.encode('utf-8')
            try:
                c = b[sequence_0]
            except KeyError:
                c = self.coef_obs
                b[sequence_0] = self.coef_obs
                if DEBUG: 'Observation probability for state, part not found:',state_i.encode('utf-8'),sequence_0
                
            d[0].append(a*c)
        if len(sequence)>1:
            for t in xrange(1, len(sequence)): # for each t > 1
                state_i = None
                sequence_t = sequence[t]
                for state_i in self.states:
                    dtemp = []
                    ftemp = []
                    for j,state_j in enumerate(self.states):
                        try:
                            a = self.state_transitions[state_j]
                        except KeyError:
                            a={}
                            self.state_transitions[state_j]=a
                            if DEBUG: print 'State transition dict for state not found:',state_j.encode('utf-8')
                        try:
                            b = a[state_i]
                        except KeyError:
                            b = self.coef_trans
                            a[state_i]=self.coef_trans
                            if DEBUG: print 'State transition for state, state not found:',state_j.encode('utf-8'),state_i.encode('utf-8')
                        
                        daji = d[t-1][j] * b
                        
                        try:
                            c = self.observations_prob[state_i]
                        except KeyError:
                            c = {}
                            self.observations_prob[state_i] = c
                            if DEBUG: assert self.observations_prob[state_i] == {}
                            if DEBUG: 'Observation for state not found:', state_i.encode('utf-8')                       
                        try:
                            e = c[sequence_t]
                        except KeyError:
                            e = self.coef_obs
                            c[sequence_t]=self.coef_obs
                            if DEBUG: 'Observation probability for state, part not found:',state_i.encode('utf-8'),sequence_t                       
                            
                        dtemp.append(daji * e)
                        ftemp.append(daji)
                    # print 'd',
                    # pprint.pprint(dtemp)
                    # print 'f',
                    # pprint.pprint(ftemp)
        
                    f[t].append(ftemp.index(max(ftemp)))
                    d[t].append(max(dtemp))
                # print '-'*20
                
        else:
            t=0
        # print 'final f:'    
        # pprint.pprint(f)
        # print 'final d:'
        # pprint.pprint(d) 
        state_seq = []
        it = d[t].index(max(d[t]))
        state_seq.append(self.states[it])
        for t in reversed(range(0, len(sequence)-1)):
            it = f[t+1][it]
            state_seq.append(self.states[it])
                          
        state_seq.reverse()
        return state_seq
                
    
          
    
    def forward(self, sequence):
        obs_prob_get = self.observations_prob.get

        a = []
        for i in range(0, len(sequence)): # initialize a
            a.append([])
        sequence_0 = sequence[0]
        for state_i in self.states: # for t = 1
            try:
                start_prob_i = self.start_probabilities[state_i]
            except KeyError:
                start_prob_i = self.coef_trans
                self.start_probabilities[state_i] = start_prob_i
            try:
                obs_prob_i = self.observations_prob[state_i]
            except KeyError:
                obs_prob_i = {}
                self.observations_prob[state_i] = obs_prob_i
            try:
                obs_prob_i_0 = obs_prob_i[sequence_0]
            except KeyError:
                obs_prob_i_0 = self.coef_obs
                obs_prob_i[sequence_0] = obs_prob_i_0
            a[0].append(start_prob_i * obs_prob_i_0)
        for t in xrange(1, len(sequence)): # for each t > 1
            sequence_t=sequence[t]
            for state_j in self.states: # for each state
                suma = 0
                for i, state_i in enumerate(self.states):
                    ati = a[t-1][i]
                    try:
                        state_trans_i = self.state_transitions[state_i]
                    except KeyError:
                        state_trans_i = {}
                        self.state_transitions[state_i] = state_trans_i
                    try:
                        state_trans_i_j = state_trans_i[state_j]
                    except KeyError:
                        state_trans_i_j = self.coef_trans    
                        state_trans_i[state_j] = state_trans_i_j
                    aij = state_trans_i_j
                    suma=suma + ati*aij# SUM[at(i)*aij]
                try:
                    obs_prob_j = self.observations_prob[state_j]
                except KeyError:
                    obs_prob_j = {}
                    self.observations_prob[state_j] = obs_prob_j
                try:
                    obs_prob_j_t = obs_prob_j[sequence_t]
                except KeyError:
                    obs_prob_j_t = self.coef_obs
                    obs_prob_j[sequence_t] = obs_prob_j_t
                bjkt = obs_prob_j_t
                a[t].append(bjkt * suma)
        total_prob = 0
        t = len(sequence)-1
        for j in xrange(0, self.len_states):
            total_prob+= a[t][j]
        return total_prob
    
    


def test_numpy(self):
    
    import numpy
    A = numpy.array([[0.5, 0.25, 0.25],
    [0.375,0.125,0.375],
    [0.125,0.675,0.375]])
    Pi = numpy.array([0.63,0.17,0.20])
    B = numpy.array([[0.60,0.20,0.15,0.05],
    [0.25,0.25,0.25,0.25],
    [0.05,0.10,0.35,0.50]])
    obs = [0,2,3] # dry, damp, soggy
    
    d = numpy.zeros((3,3),float)
    f = numpy.zeros((3,3),int)